package com.coolw.log4j2.mask.log;

import com.google.common.collect.Maps;
import lombok.extern.slf4j.Slf4j;
import org.apache.logging.log4j.core.Layout;
import org.apache.logging.log4j.core.LogEvent;
import org.apache.logging.log4j.core.config.Node;
import org.apache.logging.log4j.core.config.plugins.Plugin;
import org.apache.logging.log4j.core.config.plugins.PluginAttribute;
import org.apache.logging.log4j.core.config.plugins.PluginFactory;
import org.apache.logging.log4j.core.layout.AbstractStringLayout;
import org.apache.logging.log4j.core.layout.PatternLayout;
import org.dromara.hutool.core.collection.CollUtil;
import org.dromara.hutool.core.text.StrUtil;

import java.nio.charset.Charset;
import java.util.Arrays;
import java.util.Map;
import java.util.Set;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
import java.util.stream.Collectors;

/**
 * 自定义日志格式,参照{@link PatternLayout}
 *
 * @author coolw
 * @date 2024/4/19 17:04
 */
@Slf4j
@Plugin(name = "CustomPatternLayout", category = Node.CATEGORY, elementType = Layout.ELEMENT_TYPE, printObject = true)
public class CustomPatternLayout extends AbstractStringLayout {

    private final PatternLayout patternLayout;

    /**
     * 正则表达式对应的Pattern
     */
    private static final Map<String, Pattern> REG_PATTERN_MAP = Maps.newConcurrentMap();

    /**
     * 脱敏关键字对应的正则表达式
     */
    private static final Map<String, String> KEY_REG_MAP = Maps.newConcurrentMap();

    public CustomPatternLayout(Charset charset, String pattern) {
        super(charset);
        patternLayout = PatternLayout.newBuilder().withPattern(pattern).build();
        initRule();
    }

    private void initRule() {
        final Map<String, String> regularMap = Log4j2MaskRule.regularMap;
        if (CollUtil.isEmpty(regularMap)) {
            log.info(">>>>>> 日志脱敏规则为空");
            return;
        }

        try {
            regularMap.forEach((a, b) -> {
                if (StrUtil.isNotBlank(a)) {
                    Map<String, String> keyRegMap = Arrays.stream(a.split(","))
                            .collect(Collectors.toMap(c -> c, w -> b, (key1, key2) -> key1));
                    KEY_REG_MAP.putAll(keyRegMap);
                }
                Pattern pattern = Pattern.compile(b);
                REG_PATTERN_MAP.put(b, pattern);
            });
        } catch (Exception e) {
            log.warn(">>>>>> 初始化日志脱敏规则异常,regularMap:{}", regularMap, e);
        }
    }

    /**
     * 日志脱敏
     * 1.判断配置文件中是否已经配置需要脱敏字段
     * 2.判断内容是否有需要脱敏的敏感信息
     * 2.1 没有需要脱敏信息直接返回
     * 2.2 处理需要脱敏的字段
     */
    private String hideMarkLog(String logStr) {
        // 1.判断配置文件中是否已经配置需要脱敏字段
        if (StrUtil.isBlank(logStr) || CollUtil.isEmpty(KEY_REG_MAP) || CollUtil.isEmpty(REG_PATTERN_MAP)) {
            return logStr;
        }

        // 2.判断内容是否有需要脱敏的敏感信息
        try {
            Set<String> keys = KEY_REG_MAP.keySet();
            for (String key : keys) {
                if (logStr.contains(key)) {
                    String regExp = KEY_REG_MAP.get(key);
                    logStr = matchAndMaskHandle(logStr, regExp, key);
                }
            }
            return logStr;
        } catch (Exception e) {
            log.warn(">>>>>> 日志脱敏处理异常", e);
            return logStr;
        }
    }

    /**
     * 正则匹配对应的对象
     */
    private static String matchAndMaskHandle(String msg, String regExp, String key) {
        Pattern pattern = REG_PATTERN_MAP.get(regExp);
        if (pattern == null) {
            log.info(">>>>>> 没有匹配到对应的正则表达式,regExp:{}", regExp);
            return msg;
        }

        Matcher matcher = pattern.matcher(msg);
        int length = key.length() + 5;
        boolean contains = Log4j2MaskRule.NAME_KEY.contains(key);
        String hiddenStr;
        while (matcher.find()) {
            String originStr = matcher.group();
            if (contains) {
                // 计算关键词和需要脱敏词的距离小于5
                int i = msg.indexOf(originStr);
                if (i < 0) {
                    continue;
                }
                int span = i - length;
                int startIndex = Math.max(span, 0);
                String substring = msg.substring(startIndex, i);
                if (StrUtil.isBlank(substring) || !substring.contains(key)) {
                    continue;
                }
            }
            hiddenStr = hideMarkStr(originStr);
            msg = msg.replace(originStr, hiddenStr);
        }
        return msg;
    }

    /**
     * 标记敏感文字规则
     */
    private static String hideMarkStr(String needHideMark) {
        if (StrUtil.isBlank(needHideMark)) {
            return "";
        }

        int startSize;
        int endSize;
        int mark;
        int length = needHideMark.length();

        StringBuilder hideRegBuilder = new StringBuilder("(\\S{");
        StringBuilder replaceBuilder = new StringBuilder("$1");

        if (length > 4) {
            int i = length / 3;
            startSize = i;
            endSize = i;
        } else {
            startSize = 1;
            endSize = 0;
        }

        mark = length - startSize - endSize;
        for (int i = 0; i < mark; i++) {
            replaceBuilder.append("*");
        }

        hideRegBuilder.append(startSize).append("})\\S*(\\S{").append(endSize).append("})");
        replaceBuilder.append("$2");
        needHideMark = needHideMark.replaceAll(hideRegBuilder.toString(), replaceBuilder.toString());
        return needHideMark;
    }

    /**
     * 创建插件
     */
    @PluginFactory
    public static Layout createLayout(@PluginAttribute(value = "pattern") final String pattern,
                                      @PluginAttribute(value = "charset") final Charset charset) {
        return new CustomPatternLayout(charset, pattern);
    }

    @Override
    public String toSerializable(LogEvent event) {
        return hideMarkLog(patternLayout.toSerializable(event));
    }
}
